"""
音频录制模块
使用sounddevice录制用户语音
"""
import sounddevice as sd
import soundfile as sf
import numpy as np
import threading
import time
from typing import Optional, Callable
import os

class AudioRecorder:
    def __init__(self, sample_rate: int = 16000, channels: int = 1):
        """
        初始化录音器
        
        Args:
            sample_rate: 采样率 (Hz)
            channels: 声道数
        """
        self.sample_rate = sample_rate
        self.channels = channels
        self.is_recording = False
        self.recorded_data = None
        self.recording_thread = None
        
        # 回调函数
        self.on_recording_start: Optional[Callable] = None
        self.on_recording_stop: Optional[Callable] = None
        self.on_recording_update: Optional[Callable] = None
    
    def get_audio_devices(self) -> dict:
        """
        获取可用的音频设备
        
        Returns:
            包含输入和输出设备的字典
        """
        devices = sd.query_devices()
        input_devices = []
        output_devices = []
        
        for i, device in enumerate(devices):
            device_info = {
                'index': i,
                'name': device['name'],
                'channels': device['max_input_channels'] if device['max_input_channels'] > 0 else device['max_output_channels'],
                'sample_rate': device['default_samplerate']
            }
            
            if device['max_input_channels'] > 0:
                input_devices.append(device_info)
            if device['max_output_channels'] > 0:
                output_devices.append(device_info)
        
        return {
            'input_devices': input_devices,
            'output_devices': output_devices,
            'default_input': sd.default.device[0],
            'default_output': sd.default.device[1]
        }
    
    def set_device(self, device_index: Optional[int] = None):
        """
        设置录音设备
        
        Args:
            device_index: 设备索引，None表示使用默认设备
        """
        if device_index is not None:
            sd.default.device[0] = device_index
    
    def start_recording(self, max_duration: float = 10.0):
        """
        开始录音
        
        Args:
            max_duration: 最大录音时长（秒）
        """
        if self.is_recording:
            print("已在录音中...")
            return
        
        self.is_recording = True
        self.recorded_data = []
        
        # 触发开始录音回调
        if self.on_recording_start:
            self.on_recording_start()
        
        def recording_worker():
            try:
                # 计算缓冲区大小（100ms）
                blocksize = int(self.sample_rate * 0.1)
                
                with sd.InputStream(
                    samplerate=self.sample_rate,
                    channels=self.channels,
                    blocksize=blocksize,
                    dtype=np.float32
                ) as stream:
                    
                    start_time = time.time()
                    
                    while self.is_recording and (time.time() - start_time) < max_duration:
                        # 读取音频数据
                        data, overflowed = stream.read(blocksize)
                        
                        if overflowed:
                            print("音频缓冲区溢出")
                        
                        # 存储数据
                        self.recorded_data.append(data.copy())
                        
                        # 触发更新回调
                        if self.on_recording_update:
                            duration = time.time() - start_time
                            self.on_recording_update(duration, len(self.recorded_data))
                        
                        time.sleep(0.01)  # 短暂休眠
                
                # 录音结束
                self.is_recording = False
                
                # 触发停止录音回调
                if self.on_recording_stop:
                    duration = time.time() - start_time
                    self.on_recording_stop(duration)
                    
            except Exception as e:
                print(f"录音错误: {e}")
                self.is_recording = False
                if self.on_recording_stop:
                    self.on_recording_stop(0)
        
        # 在单独线程中进行录音
        self.recording_thread = threading.Thread(target=recording_worker)
        self.recording_thread.daemon = True
        self.recording_thread.start()
    
    def stop_recording(self):
        """停止录音"""
        if self.is_recording:
            self.is_recording = False
            
            # 等待录音线程结束
            if self.recording_thread and self.recording_thread.is_alive():
                self.recording_thread.join(timeout=1.0)
    
    def get_recorded_audio(self) -> Optional[np.ndarray]:
        """
        获取录制的音频数据
        
        Returns:
            录制的音频数组，如果没有数据则返回None
        """
        if not self.recorded_data:
            return None
        
        # 合并所有音频块
        audio_data = np.concatenate(self.recorded_data, axis=0)
        
        # 如果是多声道，取第一个声道
        if len(audio_data.shape) > 1:
            audio_data = audio_data[:, 0]
        
        return audio_data
    
    def save_recording(self, file_path: str, audio_data: Optional[np.ndarray] = None) -> bool:
        """
        保存录音到文件
        
        Args:
            file_path: 保存路径
            audio_data: 音频数据，如果为None则使用最近录制的数据
            
        Returns:
            是否保存成功
        """
        try:
            if audio_data is None:
                audio_data = self.get_recorded_audio()
            
            if audio_data is None:
                print("没有可保存的音频数据")
                return False
            
            # 确保目录存在
            os.makedirs(os.path.dirname(file_path), exist_ok=True)
            
            # 保存为WAV格式
            sf.write(file_path, audio_data, self.sample_rate)
            print(f"音频已保存至: {file_path}")
            return True
            
        except Exception as e:
            print(f"保存音频失败: {e}")
            return False
    
    def play_audio(self, audio_data: np.ndarray):
        """
        播放音频数据
        
        Args:
            audio_data: 要播放的音频数据
        """
        try:
            sd.play(audio_data, self.sample_rate)
        except Exception as e:
            print(f"播放音频失败: {e}")
    
    def play_file(self, file_path: str):
        """
        播放音频文件
        
        Args:
            file_path: 音频文件路径
        """
        try:
            audio_data, sample_rate = sf.read(file_path)
            sd.play(audio_data, sample_rate)
        except Exception as e:
            print(f"播放文件失败: {e}")
    
    def stop_playback(self):
        """停止音频播放"""
        sd.stop()
    
    def get_audio_level(self, audio_data: np.ndarray) -> float:
        """
        计算音频音量等级
        
        Args:
            audio_data: 音频数据
            
        Returns:
            音量等级 (0-1)
        """
        if len(audio_data) == 0:
            return 0.0
        
        # 计算RMS（均方根）
        rms = np.sqrt(np.mean(audio_data ** 2))
        
        # 转换为分贝并标准化到0-1范围
        if rms > 0:
            db = 20 * np.log10(rms)
            # 假设-60dB到0dB的范围
            level = max(0, min(1, (db + 60) / 60))
        else:
            level = 0
        
        return level
    
    def detect_voice_activity(self, audio_data: np.ndarray, threshold: float = 0.01) -> bool:
        """
        检测是否有语音活动
        
        Args:
            audio_data: 音频数据
            threshold: 检测阈值
            
        Returns:
            是否检测到语音
        """
        level = self.get_audio_level(audio_data)
        return level > threshold

# 使用示例
if __name__ == "__main__":
    recorder = AudioRecorder()
    
    # 获取音频设备信息
    devices = recorder.get_audio_devices()
    print("可用输入设备:")
    for device in devices['input_devices']:
        print(f"  {device['index']}: {device['name']} ({device['channels']}ch, {device['sample_rate']}Hz)")
    
    print(f"\n当前默认输入设备: {devices['default_input']}")
    
    # 测试录音功能
    def on_start():
        print("开始录音...")
    
    def on_update(duration, blocks):
        print(f"\r录音中... {duration:.1f}秒 ({blocks}块)", end="")
    
    def on_stop(duration):
        print(f"\n录音结束，总时长: {duration:.1f}秒")
    
    recorder.on_recording_start = on_start
    recorder.on_recording_update = on_update
    recorder.on_recording_stop = on_stop
    
    print("\n按Enter开始录音...")
    input()
    
    recorder.start_recording(max_duration=5.0)
    
    print("按Enter停止录音...")
    input()
    
    recorder.stop_recording()
    
    # 保存录音
    audio_data = recorder.get_recorded_audio()
    if audio_data is not None:
        print(f"录制了 {len(audio_data)} 个采样点")
        print(f"音频长度: {len(audio_data) / recorder.sample_rate:.2f} 秒")
        print(f"音量等级: {recorder.get_audio_level(audio_data):.3f}")
        
        # 保存到临时文件
        temp_file = "temp_recording.wav"
        if recorder.save_recording(temp_file, audio_data):
            print(f"音频已保存到 {temp_file}")
    else:
        print("没有录制到音频数据")
